# Copyright (c) 2025 Huawei Technologies Co. Ltd.
# deepinsight is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#          http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
import argparse
import os
from pathlib import Path

from camel.types import ModelPlatformType, ModelType
from dotenv import load_dotenv
from rich.console import Console
from rich.markdown import Markdown
from rich.panel import Panel
from rich.progress import Progress
from rich.table import Table

from deepinsight.config.model import ModelConfig
from deepinsight.core.orchestrator import Orchestrator, OrchestrationResult
from deepinsight.service.deep_research import AGENT_PROCESS_TIPS_TEMPLATE
from deepinsight.utils.console_utils import display_stream

load_dotenv()

console = Console()

PROJECT_ROOT_DIR = Path(__file__).parent.parent.absolute()
MCP_CONFIG_PATH = PROJECT_ROOT_DIR / "mcp_config.json"
CONFIG_PATH = PROJECT_ROOT_DIR / "config.yaml"

def save_artifact(output_dir: Path, data: OrchestrationResult) -> None:
    """Save research artifacts to files"""
    output_dir.mkdir(exist_ok=True)

    # Example saving logic
    if data.report:
        (output_dir / 'report.md').write_text(data.report, encoding='utf-8')
        console.print(
            Panel.fit(
                f"[green]Artifacts saved to: [bold]{output_dir.absolute()}[/bold][/green]",
                title="📁 Save Complete",
                border_style="blue"
            )
        )


def get_user_query() -> str:
    """Interactive prompt for research query"""
    console.print(
        Panel.fit(
            "[bold cyan]🔍 DeepInsight Research Assistant[/bold cyan]",
            border_style="green"
        )
    )
    console.print(
        "[bold]Please enter your research question or topic[/bold] "
        "[dim](press Ctrl+C to exit):[/dim]"
    )
    while True:
        query = console.input("[bold blue]> [/bold blue]").strip()
        if query:
            return query
        console.print("[red]Please enter a valid query.[/red]")


def handle_interactive_states(orchestration: Orchestrator, query: str) -> OrchestrationResult:
    """Handle user interactions during research"""
    while True:
        result = display_stream(orchestration.run(query))
        if result.require_user_feedback:
            console.print(
                Panel.fit(
                    f"[yellow]❓ {result.require_user_feedback}[/yellow]",
                    title="📝 Additional Information Needed",
                    border_style="yellow"
                )
            )
            query = console.input("[bold cyan]💬 Your response: [/bold cyan]")

        elif result.plan_draft:
            console.print(
                Panel.fit(
                    f"[bold]{result.plan_draft}[/bold]",
                    title="📋Research Plan Draft",
                    border_style="blue"
                )
            )

            table = Table.grid(padding=1)
            table.add_column("Option", style="cyan")
            table.add_column("Description", style="white")

            table.add_row("1", "Edit draft before continuing")
            table.add_row("2", "Approve and start research")
            table.add_row("3", "Cancel research")

            console.print(table)

            choice = console.input("[bold]Select option (1-3): [/bold]").strip()


            if choice == "1":
                query = console.input("[bold]Enter your modified draft:\n[/bold]")
            elif choice == "2":
                with Progress() as progress:
                    task = progress.add_task("[cyan]Starting research...", total=1)
                    query = "开始研究"
                    progress.update(task, completed=1)
            else:
                raise KeyboardInterrupt()
        else:
            break
    return result


def main():
    parser = argparse.ArgumentParser(
        description="DeepInsight Research CLI",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "-o", "--output",
        type=Path,
        help="Directory to save research artifacts",
        default=None
    )
    parser.add_argument(
        "--verbose", "-v",
        action="count",
        default=1,
        help="Increase output verbosity"
    )

    args = parser.parse_args()

    try:
        console.print(
            Panel.fit(
                "[bold green]Initializing DeepInsight Research System[/bold green]",
                border_style="green"
            )
        )
        orchestration = Orchestrator(
            model_config=ModelConfig(
                model_platform=ModelPlatformType(os.getenv('MODEL_PLATFORM', ModelPlatformType.DEEPSEEK.value)),
                model_type=ModelType(os.getenv('MODEL_TYPE', ModelType.DEEPSEEK_CHAT.value)),
                api_key=os.getenv('MODEL_API_KEY', None),
                model_config_dict=dict(
                    stream=True
                ),
            ),
            mcp_tools_config_path=str(MCP_CONFIG_PATH),
            mcp_client_timeout=60,
            research_round_limit=1,
            execute_tips_template_dict=AGENT_PROCESS_TIPS_TEMPLATE,
        )

        # Get user query interactively
        query = get_user_query()
        console.print(
            Panel.fit(
                f"[bold]🚀 Starting research:[/bold] [cyan]{query}[/cyan]",
                border_style="cyan"
            )
        )

        # Handle interactive states
        result = handle_interactive_states(orchestration, query)

        # Final output handling
        if args.output:
            save_artifact(args.output, result)

        if result.report:
            console.print(
                Panel.fit(
                    "[bold green]🔍 Research Completed![/bold green]",
                    border_style="green"
                )
            )
            console.print(
                Panel.fit(
                    Markdown(result.report),
                    title="📄 Final Report",
                    border_style="blue"
                )
            )

    except KeyboardInterrupt:
        console.print(
            Panel.fit(
                "[red]Research cancelled by user[/red]",
                border_style="red"
            )
        )
    except Exception as e:
        console.print(
            Panel.fit(
                f"[red]❌ Research failed: {str(e)}[/red]",
                title="Error",
                border_style="red"
            )
        )
        if args.verbose > 0:
            import traceback
            traceback.print_exc()
        return 1

    return 0


if __name__ == "__main__":
    import sys

    sys.exit(main())
